package com.huawei.hiimagesegmentation.help;

import android.content.Context;
import android.content.res.AssetFileDescriptor;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.os.SystemClock;
import android.util.Log;

import androidx.core.graphics.ColorUtils;

import com.huawei.hiimagesegmentation.bean.ModelTrackingResult;

import org.tensorflow.lite.Delegate;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.gpu.GpuDelegate;

import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.util.HashSet;
import java.util.Random;

public class TrackingMobile {
    private static final String TAG = "TrackingMobile";


    private static final String IMAGESEGMENTATIONMODEL = "deeplabv3_257_mv_gpu.tflite";
    private static final int imageSize = 257;
    public static final int NUM_CLASSES = 21;
    private static final float IMAGE_MEAN = 127.5F;
    private static final float IMAGE_STD = 127.5F;

    public static final int[] segmentColors = new int[2];
    public static final String[] labelsArrays = new String[]{"background", "aeroplane", "bicycle", "bird", "boat", "bottle",
            "bus", "car", "cat", "chair", "cow",
            "dining table", "dog", "horse", "motorbike",
            "person", "potted plant", "sheep", "sofa", "train", "tv"};

    private Context context;
    private GpuDelegate gpuDelegate;
    private ByteBuffer segmentationMasks;
    private Interpreter interpreter;

    private long fullTimeExecutionTime;
    private long preprocessTime;
    private long imageSegmentationTime;
    private long maskFlatteningTime;
    private int numberThreads = 4;
    private boolean useGPU;

    private Bitmap maskBitmap;
    private Bitmap resultBitmap;
    private HashSet itemsFound;

    public TrackingMobile(Context context, boolean useGPU) {
        this.context = context;
        this.useGPU = useGPU;
        interpreter = getInterpreter();

        segmentationMasks = ByteBuffer.allocateDirect(1 * imageSize * imageSize * NUM_CLASSES * 4);
        segmentationMasks.order(ByteOrder.nativeOrder());

        Random random =new Random(System.currentTimeMillis());
        segmentColors[0] = Color.TRANSPARENT;
        segmentColors[1] = Color.RED;
//        for (int i = 1; i < NUM_CLASSES ; i++) {
//            segmentColors[i] = Color.argb(128, getRandomRGBInt(random),getRandomRGBInt(random),getRandomRGBInt(random));
//        }

    }


    private int getRandomRGBInt(Random random) {
        return (int) (255 * random.nextFloat());
    }

    private Interpreter getInterpreter() {
        Interpreter.Options tfliteOptions = new Interpreter.Options();
        tfliteOptions.setNumThreads(numberThreads);
        this.gpuDelegate = (GpuDelegate) null;
        if (useGPU) {
            this.gpuDelegate = new GpuDelegate();
            tfliteOptions.addDelegate((Delegate) this.gpuDelegate);
        }

        return new Interpreter((ByteBuffer) loadModelFile(), tfliteOptions);
    }

    private MappedByteBuffer loadModelFile() {
        MappedByteBuffer retFile = null;
        try {
            AssetFileDescriptor fileDescriptor = context.getAssets().openFd(IMAGESEGMENTATIONMODEL);
            FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor());
            FileChannel fileChannel = inputStream.getChannel();
            long startOffset = fileDescriptor.getStartOffset();
            long declaredLength = fileDescriptor.getDeclaredLength();
            retFile = fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
            fileDescriptor.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return retFile;
    }

    public ModelTrackingResult execute(Bitmap bitmap) {
        this.fullTimeExecutionTime = SystemClock.uptimeMillis();
        this.preprocessTime = SystemClock.uptimeMillis();

        Bitmap scaledBitmap = ImageUtils.scaleBitmapAndKeepRatio(bitmap, imageSize, imageSize);
        ByteBuffer contentArray = ImageUtils.bitmapToByteBuffer(scaledBitmap, imageSize, imageSize, IMAGE_MEAN, IMAGE_STD);

        this.preprocessTime = SystemClock.uptimeMillis() - this.preprocessTime;
        this.imageSegmentationTime = SystemClock.uptimeMillis();
        this.interpreter.run(contentArray, this.segmentationMasks);
        this.imageSegmentationTime = SystemClock.uptimeMillis() - this.imageSegmentationTime;
        Log.d(TAG, "Time to run the model " + this.imageSegmentationTime);
        this.maskFlatteningTime = SystemClock.uptimeMillis();
        convertBytebufferMaskToBitmap(
                segmentationMasks, imageSize, imageSize, scaledBitmap,
                segmentColors
        );
        this.maskFlatteningTime = SystemClock.uptimeMillis() - this.maskFlatteningTime;
        Log.d(TAG, "Time to flatten the mask result " + this.maskFlatteningTime);
        this.fullTimeExecutionTime = SystemClock.uptimeMillis() - this.fullTimeExecutionTime;
        Log.d(TAG, "Total time execution " + this.fullTimeExecutionTime);
        return new ModelTrackingResult(resultBitmap, scaledBitmap, maskBitmap, this.formatExecutionLog(), itemsFound);
    }


    private void convertBytebufferMaskToBitmap(ByteBuffer inputBuffer, int imageWidth,
                                               int imageHeight, Bitmap backgroundImage, int[] colors) {
        Bitmap.Config conf = Bitmap.Config.ARGB_8888;
        maskBitmap = Bitmap.createBitmap(imageWidth, imageHeight, conf);
        resultBitmap = Bitmap.createBitmap(imageWidth, imageHeight, conf);
        Bitmap scaledBackgroundImage =
                ImageUtils.scaleBitmapAndKeepRatio(backgroundImage, imageWidth, imageHeight);

        int[][] mSegmentBits = new int[imageWidth][imageHeight];
        itemsFound = new HashSet();
        inputBuffer.rewind();

        for (int y = 0; y < imageHeight; y++) {
            for (int x = 0; x < imageWidth; x++) {
                float maxVal = 0f;
                mSegmentBits[x][y] = 0;

                for (int i = 0; i < NUM_CLASSES; i++) {
                    float value = inputBuffer.getFloat(
                            (y * imageWidth * NUM_CLASSES + x * NUM_CLASSES + i) * 4);

                    if (i == 0 || value > maxVal) {
                        maxVal = value;
                        if (i == 15) {
                            mSegmentBits[x][y] = i;
                        }else{
                            mSegmentBits[x][y] = 0;
                        }
                    }
                }

                itemsFound.add(mSegmentBits[x][y]);


                int newPixelColor = ColorUtils.compositeColors(
                        colors[mSegmentBits[x][y] == 0 ? 0 : 1],
                        scaledBackgroundImage.getPixel(x, y)
                );

                resultBitmap.setPixel(x, y, newPixelColor);
                maskBitmap.setPixel(x, y, mSegmentBits[x][y] == 0 ? colors[0] : scaledBackgroundImage.getPixel(x, y));
            }
        }
    }


    private final String formatExecutionLog() {
        StringBuilder sb = new StringBuilder();
        sb.append("Input Image Size: " + imageSize * imageSize + '\n');
//        sb.append("GPU enabled: " + this.useGPU + '\n');
//        sb.append("Number of threads: " + this.numberThreads + '\n');
        sb.append("Pre-process execution time: " + this.preprocessTime + " ms\n");
        sb.append("Model execution time: " + this.imageSegmentationTime + " ms\n");
        sb.append("Mask flatten time: " + this.maskFlatteningTime + " ms\n");
        sb.append("Full execution time: " + this.fullTimeExecutionTime + " ms\n");
        return sb.toString();
    }

}
